from math import log10
from langchain_community.vectorstores import Chroma
from typing import List
import os
import logging

class VecStore:
    def __init__(self, embeddings, persist_dir="vectorstore/chroma"):
        self.embeddings = embeddings
        self.persist_dir = persist_dir
        self.vectorstore = None
        self.logger = logging.getLogger(__name__)

    def initialize(self, docs_dir="local_knowledge"):
        if not os.path.exists(os.path.join(self.persist_dir, "chroma-collections.parquet")):
            self.logger.info("正在初始化本地知识库...")
            try:
                from langchain_community.document_loaders import DirectoryLoader
                from langchain_text_splitters import RecursiveCharacterTextSplitter

                loader = DirectoryLoader(docs_dir, glob="**/*.txt")
                documents = loader.load()

                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=1000,
                    chunk_overlap=200
                )
                splits = text_splitter.split_documents(documents)

                self.vectorstore = Chroma.from_documents(
                    documents=splits,
                    embedding=self.embeddings,
                    persist_directory=self.persist_dir
                )
                self.vectorstore.persist()
                self.logger.info(f"知识库已初始化到 {self.persist_dir}")
            except Exception as e:
                self.logger.error(f"知识库初始化失败: {str(e)}")
                raise

    def load(self):
        try:
            self.vectorstore = Chroma(
                persist_directory=self.persist_dir,
                embedding_function=self.embeddings
            )
            return self.vectorstore
        except Exception as e:
            self.logger.error(f"向量数据库加载失败: {str(e)}")
            raise

    def similarity_search(self, query: str, k: int = 3) -> List[str]:
        if not self.vectorstore:
            self.logger.error("向量存储未初始化，请先调用load()或initialize()方法")
            raise ValueError("Vector store not initialized")
        return self.vectorstore.similarity_search(query, k=k)

def get_context_and_sources(user_message:str, vec_store:VecStore, MAX_CONTEXT_LENGTH, logger:logging):
    """
    执行Chroma的相似度搜索，返回上下文和源信息

    参数:
    user_message (str): 用户输入的消息
    vec_store: Chroma向量存储对象
    MAX_CONTEXT_LENGTH (int): 上下文最大长度
    logger: 日志记录器对象

    返回:
    tuple: (context, sources)，context为拼接的上下文内容，sources为源信息列表
    """
    try:
        vec_store.load()
        # 修改8：Chroma的相似度搜索
        docs = vec_store.similarity_search(user_message, k=3)
        context = "\n".join([d.page_content for d in docs])[:MAX_CONTEXT_LENGTH]
        sources = [{
            "title": os.path.basename(d.metadata["source"]), # 获取源文件
            "excerpt": d.page_content[:100]+"..."
        } for d in docs]
        return context, sources
    except Exception as e:
        logger.error(f"知识库检索失败: {str(e)}")
        return "", []